import copy
import json
import os

from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

from baseconf import BASE_DISK

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

model_path="/model_path/mdeberta-v3-base-squad2" if os.path.exists(
                        "/etc") else BASE_DISK + ":/model_path/mdeberta-v3-base-squad2"

en_match_words = {
"scerario" : ["suburbs","city street","expressway","tunnel","parking-lot","gas or charging stations","unknown","road in valley"],
"weather" : ["clear","cloudy","raining","foggy","snowy","unknown","snowed"],
"period" : ["daytime","dawn or dusk","night","unknown"],
"road_structure" : ["normal","crossroads","T-junction","ramp","lane merging","parking lot entrance","round about","unknown"],
"general_obstacle" : ["nothing","speed bumper","traffic cone","water horse","stone","manhole cover","nothing","unknown"],
"abnormal_condition" : ["uneven","oil or water stain","standing water","cracked","nothing","unknown","snow cover"],
"ego_car_behavior" : ["slow down","go straight","turn right","turn left","stop","U-turn","speed up","lane change","others","braking"],
"closest_participants_type" : ["passenger car","bus","truck","pedestrain","policeman","nothing","others","unknown"],
"closest_participants_behavior" : ["slow down","go straight","turn right","turn left","stop","U-turn","speed up","lane change","others","braking","nothing"],
}

en_match_prompteverywords = {
"scerario" : ["The car was driving in the suburbs",
              "The car is driving in the city",
              "The car is on the highway",
              "The car was driving through the tunnel",
              "Cars in the parking lot",
              "The car is at a gas station or charging station",
              "unknown",
              "The car was driving through the valley"],
"weather" : ["The weather is sunny",
             "It is a cloudy day",
             "It is snowing now",
             "It is a foggy day",
             "There's snow outside",
             "unknown",
             "It has snowed"],
"period" : ["It is daylight",
            "It's dawn or dusk"
            ,"It is night"
            ,"unknown"],
"road_structure" : ["Driving on normal roads",
                    "Driving on intersection roads",
                    "Driving on T-junction",
                    "Cars enter or leave the ramp",
                    "The car is entering the lane into the road",
                    "Cars enter and exit the parking lot",
                    "The car is driving on the roundabout road",
                    "unknown"],
"general_obstacle" : ["Nothing on the road" ,
                      "There are speed bumps  in the middle of the road",
                      "There are traffic cone  in the middle of the road",
                      "There are  water horse  in the middle of the road",
                      "There are rocks or pieces  in the middle of the road",
                      "There are a manhole cover  in the middle of the road",
                      "Nothing in the middle of the road",
                      "unknown"],
"abnormal_condition" : ["The road is uneven and bumpy",
                        "There are oil or water stains on the road",
                        "There is standing water in the road",
                        "There are some cracks in the road",
                        "Nothing on the road",
                        "unknown",
                        "There is snow on the road"],
"ego_car_behavior" : ["The video car finally slow down",
                      "the video car is going straight",
                      "the video car is turning right",
                      "the video car is turning left",
                      "The video car finally stopped",
                      "The video car finally turned around",
                      "This video car is accelerating",
                      "This video car is changing lanes",
                      "others",
                      "braking"],
"closest_participants_type" : ["In front is the back or front of the car",
                               "In front is the back or front of the  bus",
                               "In front is the back or front of the  truck",
                               "The nearest person in front of the video is one or more ordinary pedestrians",
                               "The closest thing in front of the video is one or more police officers",
                               "There's nothing in front of the video but the road",
                               "others",
                               "unknown"],
"closest_participants_behavior" : ["A pedestrian or vehicle in front of you is slowing down",
                                   "Pedestrians or vehicles in front of you keep driving or walking",
                                   "The pedestrian or vehicle in front is turning right",
                                   "The pedestrian or vehicle in front is turning left",
                                   "Complete stop of pedestrians or vehicles ahead",
                                   "A pedestrian or vehicle in front is making a U-turn",
                                   "The pedestrian or vehicle in front is accelerating forward",
                                   "A pedestrian or vehicle in front of you is changing lanes",
                                   "others",
                                    "The pedestrian or vehicle in front of him is braking or the taillight is red but he is not stopping",
                                    "nothing"]
}

classifier = pipeline("zero-shot-classification", model=model_path,device=device)
sequence_to_classify = "The closetest participant exits the car and walks around the front of the car to the passenger side, where he opens the trunk and pulls out a box, which he carries back to the car."
candidate_labels = ["Nothing on the road" ,
                      "There are speed bumps  in the middle of the road",
                      "There are traffic cone  in the middle of the road",
                      "There are  water horse  in the middle of the road",
                      "There are rocks or pieces  in the middle of the road",
                      "There are a manhole cover  in the middle of the road",
                      "Nothing in the middle of the road",
                      "unknown"]
    # ["passenger car","bus","truck","pedestrain","policeman","nothing","others","unknown"]

with open("./autofrist_result_pre.json", 'r') as f:
    parsed_data = json.load(f)


# submit_json_bert = copy.deepcopy(parsed_data)
# for single_result in submit_json_bert["test_results"]:
#     for key in single_result.keys():
#         if key in ["clip_id", "scerario", "weather", "period", "road_structure"]:
#             continue
#         desc = single_result[key]
#         # candidate_labels=en_match_prompteverywords[key]
#         candidate_labels=en_match_words[key]
#         # 首先使用直接判断
#         if key=="general_obstacle":
#             pass
#         elif key=="abnormal_condition":
#             pass
#         elif key=="ego_car_behavior":
#             pass
#         elif key=="closest_participants_type":
#             pass
#         elif key=="closest_participants_behavior":
#             pass
#
#
#         output = classifier(desc, candidate_labels, multi_label=False)
#         maxindex = output['scores'].index(max(output['scores']))
#         print(f"{desc}:{key}:{en_match_words[key][maxindex]}")
#         single_result[key] = en_match_words[key][maxindex]
#
#
#
# with open(file="./autofrist_result2.json",encoding="utf-8",mode="w") as f:
#     json_data = json.dumps(submit_json_bert)
#     f.write(json_data)





tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path,)
model.to(device)
premise = "The video describes a few obstacles, including some trees, but for the most part, the road is clear and the car is able to drive freely."
hypothesis = "The road is clear and free of obstacles.'"
# hypothesis = "There are no obstacles in front of the car and it is free to drive"
# hypothesis = "There are no obstacles in front of the car and it is free to drive"
# hypothesis = "There are no obstacles in front of the car and it is free to drive"

input = tokenizer(premise, hypothesis, truncation=True, return_tensors="pt")
output = model(input["input_ids"].to(device))  # device = "cuda:0" or "cpu"
prediction = torch.softmax(output["logits"][0], -1).tolist()
label_names = ["entailment", "neutral", "contradiction"]
prediction = {name: round(float(pred) * 100, 1) for pred, name in zip(prediction, label_names)}
print(prediction)






